{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bda934fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0dcf7d63",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10, 3, 5])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a=torch.randn(10,3,4)\n",
    "b = torch.randn(10,4,5)\n",
    "torch.bmm(a,b).shape\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cade7ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "Q = torch.randn(1, 1, 32)\n",
    "K = torch.randn(1, 1, 32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ee79d24c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 64])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.concat((Q[0],K[0]),1).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "745bc0d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "l = nn.Linear(32+32,32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "2710a1e0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 32])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l(torch.concat((Q[0],K[0]),1)).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "bdf20419",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.1129, -0.1785,  0.4034,  0.1109,  0.1229, -0.0745, -0.2436,  0.4765,\n",
       "          0.8940,  0.5423, -1.0626, -0.9182, -0.1148, -0.3494, -0.0559,  0.2820,\n",
       "         -0.2363, -0.3838,  0.4071, -0.6858,  0.3512, -0.1272,  1.5519,  0.0885,\n",
       "          0.7725, -0.4922,  0.2253,  0.3613, -0.9058,  1.1619,  0.9938,  0.2784]],\n",
       "       grad_fn=<AddmmBackward0>)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l(torch.concat((Q[0],K[0]),1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "563eef28",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.0264, 0.0197, 0.0352, 0.0263, 0.0266, 0.0219, 0.0185, 0.0379, 0.0576,\n",
       "         0.0405, 0.0081, 0.0094, 0.0210, 0.0166, 0.0223, 0.0312, 0.0186, 0.0160,\n",
       "         0.0354, 0.0119, 0.0334, 0.0207, 0.1111, 0.0257, 0.0510, 0.0144, 0.0295,\n",
       "         0.0338, 0.0095, 0.0752, 0.0636, 0.0311]], grad_fn=<SoftmaxBackward0>)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.softmax(l(torch.concat((Q[0],K[0]),1)),axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "6e79839f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 3, 7],\n",
       "       [4, 1, 5],\n",
       "       [3, 4, 3]])"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 假设 preds = [0.2, 0.3, 0.5]\n",
    "np.random.multinomial(10, [0.2, 0.3, 0.5], 3)  # 可能输出 [[0, 1, 0]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7466d2b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch_py38",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
